package cc.rengu.oltp.utility.util.sm2;

import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.math.BigInteger;
import java.security.InvalidKeyException;
import java.security.SecureRandom;
import java.util.Arrays;

import cc.rengu.oltp.utility.util.SM2Util;
import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.Digest;
import org.bouncycastle.crypto.DigestDerivationFunction;
import org.bouncycastle.crypto.InvalidCipherTextException;
import org.bouncycastle.crypto.digests.SM3Digest;
import org.bouncycastle.crypto.generators.ECKeyPairGenerator;
import org.bouncycastle.crypto.generators.KDF2BytesGenerator;
import org.bouncycastle.crypto.params.ECKeyGenerationParameters;
import org.bouncycastle.crypto.params.ECPrivateKeyParameters;
import org.bouncycastle.crypto.params.ECPublicKeyParameters;
import org.bouncycastle.crypto.params.KDFParameters;
import org.bouncycastle.crypto.params.ParametersWithRandom;


/**
 * SM2 非对称、公钥加密,签名验签：基于椭圆曲线
 */
public class Sm2Engine {


    // 密钥派生函数
    private final DigestDerivationFunction kdf;
    // 杂凑函数、摘要/哈希算法
    private final Digest digest;
    private final Sm2BasicAgreement basicAgreement;

    private Boolean forEncryption;
    private ECPublicKeyParameters publicKey;
    private ECPrivateKeyParameters privateKey;
    private SecureRandom secureRandom;

    public Sm2Engine() {
        this.kdf = new KDF2BytesGenerator(new SM3Digest());
        this.digest = new SM3Digest();
        this.basicAgreement = new Sm2BasicAgreement();
    }

    public Sm2Engine(DigestDerivationFunction kdf, Digest digest) {
        this.kdf = kdf;
        this.digest = digest;
        this.basicAgreement = new Sm2BasicAgreement();
    }

    public void init(boolean forEncryption, CipherParameters parameters) throws InvalidKeyException {
        ECPrivateKeyParameters priKey = null;
        ECPublicKeyParameters pubKey = null;
        SecureRandom sr = null;

        this.forEncryption = forEncryption;
        if (forEncryption) {//加密
            if (parameters instanceof ParametersWithRandom) {
                ParametersWithRandom rParam = (ParametersWithRandom) parameters;
                pubKey = (ECPublicKeyParameters) rParam.getParameters();
                sr = rParam.getRandom();
            } else if (parameters instanceof ECPublicKeyParameters) {
                pubKey = (ECPublicKeyParameters) parameters;
                sr = new SecureRandom();
            } else {
                throw new InvalidKeyException("EC public key required for encryption");
            }
        } else {//解密
            if (parameters instanceof ParametersWithRandom) {
                ParametersWithRandom rParam = (ParametersWithRandom) parameters;
                priKey = (ECPrivateKeyParameters) rParam.getParameters();
                sr = rParam.getRandom();
            } else if (parameters instanceof ECPrivateKeyParameters) {
                priKey = (ECPrivateKeyParameters) parameters;
                sr = new SecureRandom();
            } else {
                throw new InvalidKeyException("EC private key required for decryption");
            }

        }
        this.publicKey = pubKey;
        this.privateKey = priKey;
        this.secureRandom = sr;
    }

    public byte[] processBlock(byte[] input, int inOff, int inLen) throws Exception {
        if (this.forEncryption) {
            return encryptBlock(input, inOff, inLen);
        } else {
            return DecryptBlock(input, inOff, inLen);
        }
    }

    /**
     * SM2 加密
     *
     * @param input 明文
     * @param inOff 位移
     * @param inLen 明文长度
     * @return 密文
     * @throws Exception
     */
    private byte[] encryptBlock(byte[] input, int inOff, int inLen) throws Exception {
        AsymmetricCipherKeyPair tempPair;
        ECKeyPairGenerator keyGen = new ECKeyPairGenerator();
        keyGen.init(new ECKeyGenerationParameters(publicKey.getParameters(), this.secureRandom));
        tempPair = keyGen.generateKeyPair();

        //C1=[k]G=(x1,y1)
        ECPublicKeyParameters C1 = (ECPublicKeyParameters) tempPair.getPublic(); // get temp's public key
        ECPrivateKeyParameters k = (ECPrivateKeyParameters) tempPair.getPrivate(); //get temp's private key
        byte[] x1 = SM2Util.bigIntegerToByteArray(C1.getQ().getAffineXCoord().toBigInteger());
        byte[] y1 = SM2Util.bigIntegerToByteArray(C1.getQ().getAffineYCoord().toBigInteger());
        this.basicAgreement.init(k);

        //点 S = [k]Pb，若S为无形远点，报错【vs.规范描述:h=k】
        BigInteger[] s;
        s = basicAgreement.calculateAgreement(publicKey);
        //[k]Pb=(x2,y2)
        byte[] x2 = SM2Util.bigIntegerToByteArray(s[0]);
        byte[] y2 = SM2Util.bigIntegerToByteArray(s[1]);

        //t = KDF(x2||y2, klen)
        byte[] z = new byte[x2.length + y2.length];
        System.arraycopy(x2, 0, z, 0, x2.length);
        System.arraycopy(y2, 0, z, x2.length, y2.length);
        KDFParameters kParam = new KDFParameters(z, null);
        byte[] t = KDFBytes(kParam, inLen);
        //C2=M(XOR)t
        byte[] C2 = new byte[inLen];
        for (int i = 0; i < inLen; i++) {
            C2[i] = (byte) (t[i] ^ input[inOff + i]);
        }
        //C3=Hash(x2||M||y2)
        byte[] C3 = new byte[digest.getDigestSize()];
        digest.update(x2, 0, x2.length);
        digest.update(input, inOff, inLen);
        digest.update(y2, 0, y2.length);
        digest.doFinal(C3, 0);

        //密文 C=C1||C2||C3,  无卡规范为： PC||C1||C3||C2
        byte[] PC = {(byte) 4};
        byte[] C = combineByteArray(PC, x1, y1, C3, C2);
        return C;
    }

    /**
     * SM2 解密
     *
     * @param inEnc 密文
     * @param inOff 位移
     * @param inLen 密文长度
     * @return 明文
     */
    private byte[] DecryptBlock(byte[] inEnc, int inOff, int inLen) throws Exception {
        InputStream is = new ByteArrayInputStream(inEnc);
        // 04
        Byte PC = (byte) is.read();
        if (!PC.equals((byte) 04)) {
            throw new InvalidCipherTextException("Invalid PC value.");
        }
        //C1=[k]G=(x1,y1)
        byte[] x1 = new byte[256 / 8];
        byte[] y1 = new byte[256 / 8];
        is.read(x1, 0, x1.length);
        is.read(y1, 0, y1.length);
        //验证C1是否买组椭圆曲线方程
        ECPublicKeyParameters tempPublic = SM2Util.sm2PubKeyGet(x1, y1);

        //C3=Hash(x2||M||y2)
        byte[] C3 = new byte[digest.getDigestSize()];
        is.read(C3, 0, C3.length);

        //C2=M(XOR)t
        int msgLen = inLen - 1 - x1.length - y1.length - C3.length;
        byte[] C2 = new byte[msgLen];
        is.read(C2, 0, C2.length);

        //点 S=[k]C1【加密：S = [k]Pb】，若S为无形远点，报错【vs规范描述:h=k】
        basicAgreement.init(privateKey);
        BigInteger[] s = basicAgreement.calculateAgreement(tempPublic);
        //点 [dB]C1=(x2,y2)【加密：[k]Pb=(x2,y2)】
        byte[] x2 = SM2Util.bigIntegerToByteArray(s[0]);
        byte[] y2 = SM2Util.bigIntegerToByteArray(s[1]);

        //t = KDF(x2||y2, klen)
        byte[] z = new byte[x2.length + y2.length];
        System.arraycopy(x2, 0, z, 0, x2.length);
        System.arraycopy(y2, 0, z, x2.length, y2.length);
        KDFParameters kParam = new KDFParameters(z, null);
        byte[] t = KDFBytes(kParam, inLen);

        //明文 M'=C2(XOR)t 【加密：C2=M(XOR)t】
        byte[] result = new byte[msgLen];
        for (int i = 0; i < msgLen; i++) {
            result[i] = (byte) (t[i] ^ C2[i]);
        }
        //C3=Hash(x2||M||y2)，并比较u和C3
        byte[] u = new byte[digest.getDigestSize()];
        digest.update(x2, 0, x2.length);
        digest.update(result, 0, result.length);
        digest.update(y2, 0, y2.length);
        digest.doFinal(u, 0);

        boolean check = Arrays.equals(u, C3);
        if (!check) {
            throw new InvalidCipherTextException("Invalid HASH VALUE.");
        }
        return result;
    }

    /**
     * 密钥派生函数
     *
     * @param kParam
     * @param length
     * @return
     */
    private byte[] KDFBytes(KDFParameters kParam, int length) {
        byte[] buf = new byte[length];
        kdf.init(kParam);
        kdf.generateBytes(buf, 0, buf.length);
        return buf;
    }

    /**
     * 将多个字节数组拼凑成一个字节数组
     *
     * @param args 源字节数组
     * @return 得到的字节数组
     */
    public static byte[] combineByteArray(byte[]... args) {
        if (args == null || args.length == 0) {
            return null;
        }

        int len = 0;
        for (int i = 0; i < args.length; i++) {
            if (args[i] == null) {
                continue;
            }

            len += args[i].length;
        }

        byte[] ret = new byte[len];
        int pos = 0;
        for (int i = 0; i < args.length; i++) {
            if (args[i] == null || args[i].length == 0) {
                continue;
            }
            System.arraycopy(args[i], 0, ret, pos, args[i].length);
            pos += args[i].length;
        }

        return ret;
    }

}
